{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Training a SSD MobileNet V2 for Pedestrian Detection with TensorFlow Object Detection API\n",
    "This notebook shows how to train a SSD MobileNet V2 object detector for pedestrian detection with [TensorFlow Object Detection API](https://github.com/tensorflow/models/tree/master/research/object_detection). The TensorFlow Object Detection API is a framework for training object detection models that offers a lot of flexibility.\n",
    "0. [Set up env variables](#head-0)\n",
    "1. [Download and Prepare Dataset](#head-1)\n",
    "2. [Create tfrecords](#head-2)\n",
    "3. [Download Pre-Trained Models](#head-3)\n",
    "4. [Model Training](#head-4)\n",
    "5. [Export the Model for Inference](#head-5) <br>\n",
    "5.1. [General Case](#head-5-1) <br>\n",
    "5.2. [SSD Model That Supposed to be Export to tflite](#head-5-2) <br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables <a class=anchor id=head-0></a>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%set_env DATASET_DIR=/experiment_dir/dataset\n",
    "%set_env PRETRAINED_MODELS_DIR =/experiment_dir/pretrained_models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Download and Prepare Dataset <a class=anchor id=head-1></a>\n",
    "We will be using [Oxford Town Center Dataset](https://www.robots.ox.ac.uk/ActiveVision/Research/Projects/2009bbenfold_headpose/project.html) for training our pedestrian detection. First we will download the video file and annotatoin csv file with `download_towncenter_video_and_labels.sh` script and then extract frames from video with [ffmpeg](https://ffmpeg.org/) tool and xml annotation files with `create_tfrecord.py` script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download Oxford Town Center Dataset video file and csv annotation file\n",
    "!bash ./download_towncenter_video_and_labels.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract images from video and save in images directory with `ffmpeg`\n",
    "!mkdir $DATASET_DIR/images\n",
    "!ffmpeg -i $DATASET_DIR/TownCentreXVID.avi -q:v 1 -start_number 0 -frames:v 4501 $DATASET_DIR/images/%d.jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create xml annotation files and save in xmls directory\n",
    "!python create_xmls.py --annotation_path $DATASET_DIR/TownCentre-groundtruth.top"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create `tfrecords` <a class=anchor id=head-2></a>\n",
    "TensorFlow Object detection API only accepts `tfrecord` file format for loading and reading datasets. So we wrote a python script to create `tfrecord` files from our images and annotations. it will create two `tfrecord` files for training and evaluation set in the `output_dir`. you can specify how to split train/eval sets by setting `validation_split`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python create_tfrecord.py \\\n",
    "--data_dir $DATASET_DIR \\\n",
    "--output_dir $DATASET_DIR \\\n",
    "--label_map_path ./label_map.pbtxt \\\n",
    "--validation_split 0.25"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Download Pre-Trained Models <a class=anchor id=head-3></a>\n",
    "TensorFlow provides large number of pre-trained models which is trained on various datasets. So you don't need to start training from scratch and you can use these models as your starting point. You just need to download one of the models from [TensorFlow Object Detection Model Zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md), unzip the file and specify its path in the training config file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir $PRETRAINED_MODELS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For ssd_mobilenet_v2 and ssd_mobilenet_v2_quantized\n",
    "!wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz -O $PRETRAINED_MODELS_DIR/ssd_mobilenet_v2_coco_2018_03_29.tar.gz\n",
    "# For ssdlite_mobilenet_v2\n",
    "!wget http://download.tensorflow.org/models/object_detection/ssdlite_mobilenet_v2_coco_2018_05_09.tar.gz -O $PRETRAINED_MODELS_DIR/ssdlite_mobilenet_v2_coco_2018_05_09.tar.gz\n",
    "# For faster_rcnn_resnet50\n",
    "!wget http://download.tensorflow.org/models/object_detection/faster_rcnn_resnet50_coco_2018_01_28.tar.gz -O $PRETRAINED_MODELS_DIR/faster_rcnn_resnet50_coco_2018_01_28.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar -xzvf $PRETRAINED_MODELS_DIR/ssd_mobilenet_v2_coco_2018_03_29.tar.gz -C $PRETRAINED_MODELS_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Model Training <a class=anchor id=head-4></a>\n",
    "By specifying the training config file, the model directory which stores the training checkpoints and logs and number of training steps you can start training the model.\n",
    "note that you can track the training procedure with tensorboard by running `tensorboar --logdir [MODEL_DIR]` in training docker container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir /experiment_dir/ped_ssd_mobilenet_v2\n",
    "%set_env PIPELINE_CONFIG_PATH=/repo/training/tf_object_detection_api/pedestrian_detection/configs/ssd_mobilenet_v2_pedestrian.config\n",
    "%set_env MODEL_DIR=/experiment_dir/ped_ssd_mobilenet_v2\n",
    "%set_env NUM_TRAIN_STEPS=50000\n",
    "%set_env SAMPLE_1_OF_N_EVAL_EXAMPLES=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python /models/research/object_detection/model_main.py --pipeline_config_path=${PIPELINE_CONFIG_PATH} --model_dir=${MODEL_DIR} --num_train_steps=${NUM_TRAIN_STEPS} --sample_1_of_n_eval_examples=$SAMPLE_1_OF_N_EVAL_EXAMPLES --alsologtostderr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Export the Model for Inference <a class=anchor id=head-5></a>\n",
    "After the model has been trained you should export the trained model to a frozen graph `.pb` file so you can use it for inference.\n",
    "Note that if your model is SSD and it is supposed to be converted to `tflite`, you should use `export_tflite_ssd_graph.py` script to export the model. (section 5.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 General Case <a class=anchor id=head-5-1></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the step number which has best performance\n",
    "STEP=33387"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $MODEL_DIR/frozen_graph\n",
    "%set_env INPUT_TYPE=image_tensor\n",
    "%set_env TRAINED_CKPT_PREFIX=/experiment_dir/ped_ssd_mobilenet_v2/model.ckpt-{STEP}\n",
    "%set_env EXPORT_DIR=/experiment_dir/ped_ssd_mobilenet_v2/frozen_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python /models/research/object_detection/export_inference_graph.py \\\n",
    "    --input_type=${INPUT_TYPE} \\\n",
    "    --pipeline_config_path=${PIPELINE_CONFIG_PATH} \\\n",
    "    --trained_checkpoint_prefix=${TRAINED_CKPT_PREFIX} \\\n",
    "    --output_directory=${EXPORT_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls $MODEL_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.2 SSD Model That Supposed to be Export to `tflite` <a class=anchor id=head-5-2></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the step number which has best performance\n",
    "STEP=33387"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $MODEL_DIR/frozen_graph_tflite\n",
    "%set_env OUTPUT_DIR=/experiment_dir/ped_ssd_mobilenet_v2/frozen_graph_tflite\n",
    "%set_env TRAINED_CKPT_PREFIX=/experiment_dir/ped_ssd_mobilenet_v2/model.ckpt-{STEP}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python /models/research/object_detection/export_tflite_ssd_graph.py \\\n",
    "--pipeline_config_path=$PIPELINE_CONFIG_PATH \\\n",
    "--trained_checkpoint_prefix=$TRAINED_CKPT_PREFIX \\\n",
    "--output_directory=$OUTPUT_DIR \\\n",
    "--add_postprocessing_op=true \\\n",
    "--max_detections=50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls $MODEL_DIR/frozen_graph_tflite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In case your model is trained with quantizatoin-aware training strategy you can convert the `.pb` frozen graph to a quantized `.tflite` model with our [docker container](https://github.com/neuralet/neuralet/blob/master/training/tf_object_detection_api/tools-toco.Dockerfile)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
